package com.webank.tikv.util;

import io.grpc.netty.GrpcSslContexts;
import io.netty.handler.ssl.SslContextBuilder;

import javax.net.ssl.KeyManagerFactory;
import javax.net.ssl.TrustManagerFactory;
import java.io.File;
import java.nio.file.Files;
import java.nio.file.Paths;
import java.security.KeyStore;

public abstract class CertContext {

    private static final String PUB_KEY_INFRA = "PKIX";

    public abstract SslContextBuilder createSslContextBuilder();

    public static class JksContext extends CertContext {
        private final String keyPath;
        private final String keyPassword;
        private final String trustPath;
        private final String trustPassword;

        public JksContext(String keyPath, String keyPassword, String trustPath, String trustPassword) {
            this.keyPath = keyPath;
            this.keyPassword = keyPassword;
            this.trustPath = trustPath;
            this.trustPassword = trustPassword;
        }

        @Override
        public SslContextBuilder createSslContextBuilder() {
            SslContextBuilder builder = GrpcSslContexts.forClient();
            try {
                if (keyPath != null && keyPassword != null) {
                    KeyStore keyStore = KeyStore.getInstance("JKS");
                    keyStore.load(Files.newInputStream(Paths.get(keyPath)), keyPassword.toCharArray());
                    KeyManagerFactory keyManagerFactory =
                        KeyManagerFactory.getInstance(KeyManagerFactory.getDefaultAlgorithm());
                    keyManagerFactory.init(keyStore, keyPassword.toCharArray());
                    builder.keyManager(keyManagerFactory);
                }
                if (trustPath != null && trustPassword != null) {
                    KeyStore trustStore = KeyStore.getInstance("JKS");
                    trustStore.load(Files.newInputStream(Paths.get(trustPath)), trustPassword.toCharArray());
                    TrustManagerFactory trustManagerFactory = TrustManagerFactory.getInstance(PUB_KEY_INFRA);
                    trustManagerFactory.init(trustStore);
                    builder.trustManager(trustManagerFactory);
                }
            } catch (Exception e) {
                throw new IllegalArgumentException(e);
            }
            return builder;
        }
    }

    public static class OpenSslContext extends CertContext {
        private final String trustPath;
        private final String chainPath;
        private final String keyPath;

        public OpenSslContext(String trustPath, String chainPath, String keyPath) {
            this.trustPath = trustPath;
            this.chainPath = chainPath;
            this.keyPath = keyPath;
        }

        @Override
        public SslContextBuilder createSslContextBuilder() {
            SslContextBuilder builder = GrpcSslContexts.forClient();
            try {
                if (trustPath != null) {
                    builder.trustManager(new File(trustPath));
                }
                if (chainPath != null && keyPath != null) {
                    builder.keyManager(new File(chainPath), new File(keyPath));
                }
            } catch (Exception e) {
                throw new IllegalArgumentException(e);
            }
            return builder;
        }
    }
}
